#include "driver_entry.h"
#include <list>
#include <string>
#include "new.h"
#include "common.hpp"

std::list<DWORD, std::non_paged_alloc>* pid_protect_list  = NULL;
std::list<DWORD, std::non_paged_alloc>* pid_black_list  = NULL;
std::list<std::non_paged_wstring, std::non_paged_alloc>* service_balck_list  = NULL;
std::list<std::non_paged_string, std::non_paged_alloc>* protect_proc_list = NULL;

typedef NTSTATUS (NTAPI* NtLoadDriver)(__in PUNICODE_STRING DriverServiceName);
NtLoadDriver nt_load_driver = NULL;


NTSTATUS NTAPI NtOpenProcessHandler(__out PHANDLE ProcessHandle, __in ACCESS_MASK DesiredAccess, 
									__in POBJECT_ATTRIBUTES ObjectAttributes, __in_opt PCLIENT_ID ClientId)
{
 	PEPROCESS peb = PsGetCurrentProcess();
 	const char* pSzProcName = PsGetProcessImageFileName(peb);
 	if (0 == _stricmp(pSzProcName, "moxia.exe"))
 	{
 		PEPROCESS peb = NULL;
 		if (NT_SUCCESS(PsLookupProcessByProcessId(ClientId->UniqueProcess, &peb)))
 		{
 			const char* pSzTargetProc = PsGetProcessImageFileName(peb);
 			if (0 == _stricmp(pSzTargetProc, "Ollydbg.exe") 
				|| 0 == _stricmp(pSzTargetProc, "ce.exe") 
				|| 0 == _stricmp(pSzTargetProc, "cheatengine.exe")
				|| 0 == _stricmp(pSzTargetProc, "Cheat Engine.exe"))
 			{
				KdPrint(("apex want to open <%s>! fuck it!\r\n",  pSzTargetProc));
 				return STATUS_ACCESS_DENIED;
 			}
		}
 	}
	return NtOpenProcess(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);
}
NTSTATUS NTAPI NtLoadDriverHandler(__in PUNICODE_STRING DriverServiceName)
{
	KdPrint(("%ws\r\n", DriverServiceName->Buffer));
	for (std::list<std::non_paged_wstring, std::non_paged_alloc>::iterator it = service_balck_list->begin(); it != service_balck_list->end(); it++)
	{
		KdPrint(("%ws\r\n", it->c_str()));
	}
	if (service_balck_list->end() != std::find(service_balck_list->begin(), service_balck_list->end(), std::non_paged_wstring(DriverServiceName->Buffer)))
	{
		KdPrint(("oh~ fuck apex!\r\n"));
		return STATUS_UNSUCCESSFUL;
	}
	return nt_load_driver(DriverServiceName);
}

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING  RegistryPath)
{
	UNICODE_STRING DeviceName,Win32Device;
	PDEVICE_OBJECT DeviceObject = NULL;
	NTSTATUS status;

	RtlInitUnicodeString(&DeviceName,L"\\Device\\MoXiaHack");
	RtlInitUnicodeString(&Win32Device,L"\\DosDevices\\MoXiaHack");

	for (unsigned i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
		DriverObject->MajorFunction[i] = DriverDefaultHandler;

	DriverObject->MajorFunction[IRP_MJ_CREATE] = DriverCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverCreateClose;
	
	DriverObject->DriverUnload = DriverUnload;
	status = IoCreateDevice(DriverObject, 0, &DeviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);

	if (!NT_SUCCESS(status))
		return status;

	if (!DeviceObject)
		return STATUS_UNEXPECTED_IO_ERROR;

	DeviceObject->Flags |= DO_DIRECT_IO;
	DeviceObject->AlignmentRequirement = FILE_WORD_ALIGNMENT;
	status = IoCreateSymbolicLink(&Win32Device, &DeviceName);

	DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

	pid_protect_list = new(NonPagedPool) std::list<DWORD, std::non_paged_alloc>;
	pid_black_list = new(NonPagedPool) std::list<DWORD, std::non_paged_alloc>;
	service_balck_list = new(NonPagedPool) std::list<std::non_paged_wstring, std::non_paged_alloc>;
	protect_proc_list = new(NonPagedPool) std::list<std::non_paged_string, std::non_paged_alloc>;

	common::ntos::page_protect_off();

	DWORD index = common::ntos::GetServiceIndex("ZwOpenProcess");
	if (0xffffffff != index)
	{
		KdPrint(("NtOpenProcess index = 0x%02X\r\n", index));
		InterlockedExchange((volatile LONG *)&KeServiceDescriptorTable.ntoskrnel.KiServiceTable[index], (LONG)&NtOpenProcessHandler);
	}
	
	index = common::ntos::GetServiceIndex("ZwLoadDriver");
	if (0xffffffff != index)
	{
		service_balck_list->push_back(std::non_paged_wstring(L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\apf003"));

		KdPrint(("NtLoadDriver index = 0x%02X\r\n", index));
		nt_load_driver = (NtLoadDriver)KeServiceDescriptorTable.ntoskrnel.KiServiceTable[index];
		InterlockedExchange((volatile LONG *)&KeServiceDescriptorTable.ntoskrnel.KiServiceTable[index], (LONG)&NtLoadDriverHandler);
	}
	common::ntos::page_protect_on();

	return STATUS_SUCCESS;
}

void DriverUnload(IN PDRIVER_OBJECT DriverObject)
{
	if (NULL != pid_protect_list)
		delete pid_protect_list;
	
	if (NULL != pid_black_list)
		delete pid_black_list;

	if (NULL != service_balck_list)
		delete service_balck_list;

	if (NULL != protect_proc_list)
		delete protect_proc_list;

	common::ntos::page_protect_off();
	DWORD index = common::ntos::GetServiceIndex("ZwOpenProcess");
	if (0xffffffff != index)
	{
		UNICODE_STRING export_name = RTL_CONSTANT_STRING(L"NtOpenProcess");
		PVOID export_proc = MmGetSystemRoutineAddress(&export_name);
		InterlockedExchange((volatile LONG *)&KeServiceDescriptorTable.ntoskrnel.KiServiceTable[index], (LONG)export_proc);
	}

	index = common::ntos::GetServiceIndex("ZwLoadDriver");
	if (0xffffffff != index)
	{
		InterlockedExchange((volatile LONG *)&KeServiceDescriptorTable.ntoskrnel.KiServiceTable[index], (LONG)nt_load_driver);
	}
	common::ntos::page_protect_on();

	UNICODE_STRING Win32Device;
	RtlInitUnicodeString(&Win32Device,L"\\DosDevices\\MoXiaHack");
	IoDeleteSymbolicLink(&Win32Device);
	IoDeleteDevice(DriverObject->DeviceObject);
}

NTSTATUS DriverCreateClose(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS DriverDefaultHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	Irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}
NTSTATUS DriverIoControlHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	NTSTATUS io_status = STATUS_SUCCESS;
	SIZE_T cb_operate = 0;

	PIO_STACK_LOCATION irp_stack = IoGetCurrentIrpStackLocation(Irp);

	ULONG io_control_code = irp_stack->Parameters.DeviceIoControl.IoControlCode;
	switch (io_control_code)
	{
	case IO_CONTROL_ADD_PROTECT_PID:
		break;
	case IO_CONTROL_ADD_BLACK_PID:
		break;
	default:
		io_status = STATUS_INVALID_VARIANT;
	}

	Irp->IoStatus.Status = io_status;
	Irp->IoStatus.Information = cb_operate;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return io_status;
}